import paddle.fluid as fluid
from PIL import Image
import numpy as np
import os

# 创建执行器
place = fluid.CPUPlace()
exe = fluid.Executor(place)
exe.run(fluid.default_startup_program())

# 保存预测模型路径
save_path = 'infer_model/'
# 从模型中获取预测程序、输入数据名称列表、分类器
[infer_program, feeded_var_names, target_var] = fluid.io.load_inference_model(dirname=save_path, executor=exe)


# 预处理图片
def load_image(file):
    img = Image.open(file)
    # 统一图像大小
    img = img.resize((224, 224), Image.ANTIALIAS)
    # 转换成numpy值
    img = np.array(img).astype(np.float32)
    # 转换成CHW
    img = img.transpose((2, 0, 1))
    # 转换成BGR
    img = img[(2, 1, 0), :, :] / 255.0
    img = np.expand_dims(img, axis=0)
    return img



def ergodic_img():
    dataPath = r'D:\Paddle\studyPaddle\note16-灾害识别\test'
    print(dataPath)
    for filename in os.listdir(dataPath):
        if filename.endswith('.jpg') and not ('pre' in filename):  # 文件名中不包含'pre'字符串
            forecast(filename);



def forecast(path):
    # 获取图片数据
    img = load_image('test/' + path)
    # 执行预测
    result = exe.run(program=infer_program, feed={feeded_var_names[0]: img}, fetch_list=target_var)
    # 显示图片并输出结果最大的label
    lab = np.argsort(result)[0][0][-1]
    names = ['火灾', '泥石流','地震', '水灾', '干旱', '暴雪']
    print('图片：%a，预测结果标签为：%d， 名称为：%s， 概率为：%f' % (path,lab, names[lab], result[0][0][lab]))

ergodic_img();